import os
import random
import shutil


def train_val_move(origin_dir, target_train, target_val):
    val_percent = 0.25
    images_list = [img for img in os.listdir(origin_dir)]
    random.shuffle(images_list)
    if not os.path.exists(target_train):
        os.makedirs(target_train)
    if not os.path.exists(target_val):
        os.makedirs(target_val)
    train_images_count = int((1 - val_percent) * len(images_list))
    val_images_count = len(images_list) - train_images_count
    train_count = 0
    for train_img in range(train_images_count):
        origin_path = os.path.join(origin_dir, images_list[train_img])
        target_path = os.path.join(target_train, images_list[train_img])
        shutil.copy(origin_path, target_path)
        train_count += 1
    val_count = 0
    for val_img in range(val_images_count):
        origin_path = os.path.join(origin_dir, images_list[train_count + val_img])
        target_path = os.path.join(target_val, images_list[train_count + val_img])
        shutil.copy(origin_path, target_path)
        val_count += 1
    print(target_train + ": ", train_count)
    print(target_val + ": ", val_count)


def distrib_tarin_val(root_path):
    class_list = []
    for _, direc, _ in os.walk(root_path):
        if len(direc) != 0:
            class_list = direc
            break
    for class_name in class_list:
        source_direc = os.path.join(root_path, class_name)
        train_direc = os.path.join(root_path, "train", class_name)
        val_direc = os.path.join(root_path, "val", class_name)
        train_val_move(source_direc, train_direc, val_direc)


if __name__ == "__main__":
    dataset_dir = "/home/sw/algo-env/PP-MbileNetV2/datasets/smoking-calling"
    distrib_tarin_val(dataset_dir)
